# Copyright (c) 2025 Bytedance Ltd. and/or its affiliates. All rights reserved.

# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at

#     http://www.apache.org/licenses/LICENSE-2.0

# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import os
import dataclasses
from typing import Literal

from accelerate import Accelerator
from transformers import HfArgumentParser
from PIL import Image
import json
import itertools
import torch

from uso.flux.pipeline import USOPipeline, preprocess_ref
from transformers import SiglipVisionModel, SiglipImageProcessor
from tqdm import tqdm


def horizontal_concat(images):
    widths, heights = zip(*(img.size for img in images))

    total_width = sum(widths)
    max_height = max(heights)

    new_im = Image.new("RGB", (total_width, max_height))

    x_offset = 0
    for img in images:
        new_im.paste(img, (x_offset, 0))
        x_offset += img.size[0]

    return new_im


@dataclasses.dataclass
class InferenceArgs:
    prompt: str | None = None
    image_paths: list[str] | None = None
    eval_json_path: str | None = None
    # offload: bool = False
    offload: bool = True
    num_images_per_prompt: int = 1
    model_type: Literal["flux-dev", "flux-dev-fp8", "flux-schnell"] = "flux-dev-fp8"
    width: int = 1024
    height: int = 1024
    num_steps: int = 25
    guidance: float = 4
    seed: int = 3407
    save_path: str = "output/inference"
    only_lora: bool = True
    concat_refs: bool = False
    lora_rank: int = 128
    pe: Literal["d", "h", "w", "o"] = "d"
    content_ref: int = 512
    ckpt_path: str | None = None
    use_siglip: bool = True
    instruct_edit: bool = False
    hf_download: bool = True


def main(args: InferenceArgs):
    accelerator = Accelerator()

    # init SigLIP model
    siglip_processor = None
    siglip_model = None

    siglip_path = '/workspace/comfyui/models/clip/siglip-so400m-patch14-384'
    if args.use_siglip:
        siglip_processor = SiglipImageProcessor.from_pretrained(
            # "google/siglip-so400m-patch14-384"
            siglip_path
        )
        siglip_model = SiglipVisionModel.from_pretrained(
            # "google/siglip-so400m-patch14-384"
            siglip_path
        )
        siglip_model.eval()
        siglip_model.to(accelerator.device)
        print("SigLIP model loaded successfully")

    pipeline = USOPipeline(
        args.model_type,
        accelerator.device,
        args.offload,
        only_lora=args.only_lora,
        lora_rank=args.lora_rank,
        hf_download=args.hf_download,
    )
    if args.use_siglip and siglip_model is not None:
        pipeline.model.vision_encoder = siglip_model
        print('-----> hook siglip encoder')

    assert (
        args.prompt is not None or args.eval_json_path is not None
    ), "Please provide either prompt or eval_json_path"

    if args.eval_json_path is not None:
        with open(args.eval_json_path, "rt") as f:
            data_dicts = json.load(f)
        data_root = os.path.dirname(args.eval_json_path)
    else:
        data_root = ""
        data_dicts = [{"prompt": args.prompt, "image_paths": args.image_paths}]

    print(
        f"process: {accelerator.num_processes}/{accelerator.process_index}, \
    process images: {len(data_dicts)}/{len(data_dicts[accelerator.process_index::accelerator.num_processes])}"
    )

    data_dicts = data_dicts[accelerator.process_index :: accelerator.num_processes]

    accelerator.wait_for_everyone()
    local_task_count = len(data_dicts) * args.num_images_per_prompt
    if accelerator.is_main_process:
        progress_bar = tqdm(total=local_task_count, desc="Generating Images")

    for (i, data_dict), j in itertools.product(
        enumerate(data_dicts), range(args.num_images_per_prompt)
    ):
        ref_imgs = []
        for _, img_path in enumerate(data_dict["image_paths"]):
            if img_path != "":
                img = Image.open(os.path.join(data_root, img_path)).convert("RGB")
                ref_imgs.append(img)
            else:
                ref_imgs.append(None)
        siglip_inputs = None
        if args.use_siglip and siglip_processor is not None:
            with torch.no_grad():
                siglip_inputs = [
                    siglip_processor(img, return_tensors="pt").to(pipeline.device) 
                    for img in ref_imgs[1:] if isinstance(img, Image.Image)
                    ]

        ref_imgs_pil = [
            preprocess_ref(img, args.content_ref) for img in ref_imgs[:1] if isinstance(img, Image.Image)
        ]

        if args.instruct_edit:
            args.width, args.height = ref_imgs_pil[0].size
            args.width, args.height = args.width * (1024 / args.content_ref), args.height * (1024 / args.content_ref)
        image_gen = pipeline(
            prompt=data_dict["prompt"],
            width=args.width,
            height=args.height,
            guidance=args.guidance,
            num_steps=args.num_steps,
            seed=args.seed + j,
            ref_imgs=ref_imgs_pil,
            pe=args.pe,
            siglip_inputs=siglip_inputs,
        )
        if args.concat_refs:
            image_gen = horizontal_concat([image_gen, *ref_imgs])

        if "save_dir" in data_dict:
            config_save_path = os.path.join(args.save_path, data_dict["save_dir"] + f"_{j}.json")
            image_save_path = os.path.join(args.save_path, data_dict["save_dir"] + f"_{j}.png")
        else:
            os.makedirs(args.save_path, exist_ok=True)
            config_save_path = os.path.join(args.save_path, f"{i}_{j}.json")
            image_save_path = os.path.join(args.save_path, f"{i}_{j}.png")

        # save config and image
        os.makedirs(os.path.dirname(image_save_path), exist_ok=True)
        image_gen.save(image_save_path)
        # ensure the prompt and image_paths are saved in the config file
        args.prompt = data_dict["prompt"]
        args.image_paths = data_dict["image_paths"]
        args_dict = vars(args)
        with open(config_save_path, "w") as f:
            json.dump(args_dict, f, indent=4)

        if accelerator.is_main_process:
            progress_bar.update(1)
    if accelerator.is_main_process:
        progress_bar.close()


if __name__ == "__main__":
    parser = HfArgumentParser([InferenceArgs])
    args = parser.parse_args_into_dataclasses()[0]
    main(args)
